
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time    : 2019/12/17 上午10:57
# @Author  : fugang_le
# @Software: PyCharm
import os
import re
import sys
sys.path.append("/data/lefugang/template1")

import json
import time
from keras.optimizers import Adam
from keras.metrics import categorical_accuracy
from keras.callbacks import EarlyStopping, ModelCheckpoint, TensorBoard
from keras.losses import categorical_crossentropy, mean_squared_error,mean_absolute_error


from src.lstm.config import Config

from src.lstm.model import build_model
from src.lstm.data_helper import get_train_data, get_tokenizer_train_data
from src.lstm.tokenization import get_tokenizer, save_tokenizer

os.environ['CUDA_VISIBLE_DEVICES'] = '0'


def train():
    # 训练模型
    print('training  ......')
    start_time = time.time()
    print(json.dumps(Config.__repr__(), indent=4))
    model = build_model()
    adam = Adam(lr=0.001, decay=0.0)
    model.compile(loss=mean_squared_error, optimizer=adam, metrics=[mean_squared_error])
    model.summary()


    early_stopping = EarlyStopping(monitor='val_loss', patience=20, verbose=2)
    model_checkpoint = ModelCheckpoint(Config.model_file, save_best_only=True, save_weights_only=True)
    tensorboard = TensorBoard(log_dir=Config.model_path, update_freq=Config.tensorboard_update_freq)


    data = get_tokenizer_train_data()
    tokenizer = get_tokenizer(data)
    save_tokenizer(tokenizer)
    train_inputs, train_labels, dev_inputs, dev_labels, test_inputs, test_labels = get_train_data(tokenizer)
    print(train_inputs)

    history = model.fit([train_inputs[0], train_inputs[1]], train_labels,
                        validation_data=([dev_inputs[0], dev_inputs[1]], dev_labels),
                        # validation_split= 0.1,
                        epochs=Config.epochs,
                        batch_size=Config.batch_size,
                        shuffle=Config.shuffle,
                        callbacks=[early_stopping, model_checkpoint, tensorboard])

    evaluate_result = model.evaluate(x=[test_inputs[0], test_inputs[1]], y=test_labels) # 返回 loss 和 acc
    spend_time = str(round((time.time() - start_time) / 60, 2))

    # logger.info("train history: {}".format(history.history))
    print("evaluate_result: {}".format(evaluate_result))
    print("model training spend time: {}m".format(spend_time))

train()